import numpy as np
from common.mnist import load_mnist
from ch4.tow_layer_net import TwoLayerNet

if __name__ == '__main__':
    (x_train, t_train), (x_test, t_test) = load_mnist(normalize=True, one_hot_label=True)
    train_loss_list = []
    train_acc_list = []
    test_acc_list = []

    # 超参数
    iters_num = 10000  # 迭代的轮数
    train_size = x_train.shape[0]  # 这是什么意思？
    batch_size = 100  # 每一百个作为一个batch
    iter_per_epoch = max(train_size / batch_size, 1)

    learning_rate = 0.1  # 学习率
    network = TwoLayerNet(input_size=784, hidden_size=50, output_size=10)
    for i in range(iters_num):
        print(i)
        batch_mask = np.random.choice(train_size, batch_size)  # 在 train_size 的范围内，抽取 batch_size个，组成一个一维数组。
        x_batch = x_train[batch_mask]
        t_batch = t_train[batch_mask]
        # 以上抽取完毕
        grad = network.numerical_gradient(x_batch, t_batch)
        for key in ('W1', 'b1', 'W2', 'b2',):
            network.params[key] -= learning_rate * grad[key]
        loss = network.loss(x_batch, t_batch)
        train_loss_list.append(loss)
        # 每一轮 epoch 都记录一下精度
        if i % iter_per_epoch == 0:
            train_acc = network.accuracy(x_train, t_train)
            test_acc = network.accuracy(x_test, t_test)
            train_acc_list.append(train_acc)
            test_acc_list.append(test_acc)
